package com.agent;

import com.agent.command.ClassCommand;
import com.agent.command.MemoryCommand;
import com.agent.command.ThreadCommand;
import com.agent.enhancer.AgentParam;
import com.agent.enhancer.MyAdvice;
import com.agent.enhancer.TimingAdvice;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.matcher.ElementMatchers;

import java.lang.instrument.Instrumentation;
import java.util.Scanner;

/**
 * @author: LunBoWang
 * @description: TODO
 * @date: 2025/6/10 下午4:36
 * @version: 1.0
 */

public class AgentMain {
    //premain方法
    public static void premain(String agentArgs, Instrumentation inst) {
        //使用bytebuddy增强类
        new AgentBuilder.Default()
                //禁止byte buddy处理时修改类名
                .disableClassFormatChanges()
                //处理时使用retransform增强
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                //打印出错信息
                .with(new AgentBuilder.Listener.WithTransformationsOnly(AgentBuilder.Listener
                        .StreamWriting.toSystemOut()))
                //匹配哪些类
                .type(ElementMatchers.isAnnotatedWith(ElementMatchers.named("org.springframework.web.bind.annotation.RestController")
                        .or(ElementMatchers.isAnnotatedWith(ElementMatchers.named("org.springframework.web.bind.annotation.Controller")))))
                //增强，使用MyAdvice通知，对所有方法都进行增强
                .transform((builder, typeDescription, classLoader, module,protectionDomain)
                        -> builder.visit(Advice
                                .withCustomMapping()
                                .bind(AgentParam.class,agentArgs)
                        .to(TimingAdvice.class).on(ElementMatchers.any())))
                .installOn(inst);
    }

    public static void agentmain(String agentArgs, Instrumentation inst) {
        Scanner scanner = new Scanner(System.in);
        while (true) {
            System.out.println("\n=== Java Agent 菜单 ===");
            System.out.println("1. 打印内存信息");
            System.out.println("2. 生成堆内存快照");
            System.out.println("3. 打印线程栈信息");
            System.out.println("4. 显示类加载器");
            System.out.println("5. 查看类源码");
            System.out.println("6. 开启方法监控");
            System.out.println("0. 退出");
            System.out.print("请选择操作：");

            int choice = scanner.nextInt();
            switch (choice) {
                case 1:
                    MemoryCommand.printMemory();
                    break;
                case 2:
                    MemoryCommand.heapDump();
                    break;
                case 3:
                    ThreadCommand.printThreadInfo();
                    break;
                case 4:
                    ClassCommand.printAllClassLoader(inst);
                    break;
                case 5:
                    ClassCommand.printClassSourceCode(inst);
                    break;
                case 6:
                    ClassCommand.enhanceClass(inst);
                    break;
                case 0:
                    scanner.close();
                    return;
                default:
                    System.out.println("无效选项，请重新选择");
            }

            // 等待用户确认
            System.out.print("按回车继续...");
            new Scanner(System.in).nextLine();
        }
    }
}
